import sqlglot
from sqlglot import exp
from fastapi import FastAPI

def extract_table_names(sql: str) -> list:
    """
    从DQL语句中提取所有涉及的表名（去重）

    参数:
        sql: 输入的DQL语句（SELECT语句）

    返回:
        表名列表（包含完整路径，如'db.schema.table'）
    """
    table_names = set()  # 用集合去重

    try:
        # 解析SQL语句（支持多种方言，默认自动识别）
        parsed = sqlglot.parse_one(sql)

        # 遍历AST中所有的表引用节点
        for node in parsed.walk():
            # 匹配表节点（Table类型）
            if isinstance(node, exp.Table):
                # 获取完整表名（可能包含数据库、schema前缀）
                full_table_name = node.name
                table_names.add(full_table_name)

    except Exception as e:
        print(f"解析SQL失败: {e}")

    return sorted(table_names)  # 排序后返回列表


# 测试示例
if __name__ == "__main__":
    # 测试用的DQL语句（包含各种复杂情况）
    test_sqls = [
        # 简单查询
        "SELECT id, name FROM users;",
        # 带别名和JOIN
        "SELECT u.name, o.order_no FROM users u JOIN orders o ON u.id = o.user_id;",
        # 带数据库和schema前缀
        "SELECT * FROM db1.public.products WHERE price > 100;",
        # 包含子查询
        "SELECT * FROM (SELECT * FROM orders WHERE status = 'paid') t WHERE amount > 500;",
        # 带视图（视为表处理）
        "SELECT * FROM user_view uv JOIN departments d ON uv.dept_id = d.id;"
    ]

    # 逐个测试并打印结果
    for i, sql in enumerate(test_sqls, 1):
        print(f"测试用例 {i}:")
        print(f"SQL语句: {sql}")
        tables = extract_table_names(sql)
        print(f"提取的表名: {tables}\n")
